package org.lc.gateway.filter;

import cn.hutool.core.util.ObjUtil;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.lc.gateway.context.ContextExtraDataGenerator;
import org.lc.gateway.context.GatewayContext;
import org.lc.gateway.enums.FilterOrderEnum;
import org.lc.gateway.enums.GatewayLogTypeEnum;
import org.lc.gateway.properties.GatewayPluginProperties;
import org.lc.gateway.util.SchemeUtil;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.cloud.gateway.filter.factory.rewrite.CachedBodyOutputMessage;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.cloud.gateway.support.BodyInserterContext;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ReactiveHttpOutputMessage;
import org.springframework.http.codec.HttpMessageReader;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.util.AntPathMatcher;
import org.springframework.util.CollectionUtils;
import org.springframework.util.MultiValueMap;
import org.springframework.web.reactive.function.BodyInserter;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.io.UnsupportedEncodingException;
import java.net.URI;
import java.net.URLEncoder;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
 * 读取并缓存请求数据
 * Quoted from @see
 * Gateway Context Filter
 *
 * @author dl
 * 2023/01/29
 */
@Slf4j
@AllArgsConstructor
@SuppressWarnings("all")
public class GatewayRequestContextFilter implements GlobalFilter, Ordered {

    private GatewayPluginProperties gatewayPluginProperties;

    private ContextExtraDataGenerator contextExtraDataGenerator;

    private static final AntPathMatcher ANT_PATH_MATCHER = new AntPathMatcher();

    /**
     * default HttpMessageReader
     */
    private static final List<HttpMessageReader<?>> MESSAGE_READERS = HandlerStrategies.withDefaults().messageReaders();

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        var request = exchange.getRequest();
        var gatewayContext = new GatewayContext();
        /* 判断是否需求读取request 数据*/
        gatewayContext.setReadRequestData(shouldReadRequestInjectionData(exchange));
        /* 判断是否需求读取response 数据*/
        gatewayContext.setReadResponseData(gatewayPluginProperties.getResponseLog());
        /* 添加上下文请求头 */
        var headers = request.getHeaders();
        gatewayContext.setRequestHeaders(headers);
        /* 添加上下问请求数据 */
        if (Objects.nonNull(contextExtraDataGenerator)) {
            var gatewayContextExtraData = contextExtraDataGenerator.generateContextExtraData(exchange);
            gatewayContext.setGatewayContextExtraData(gatewayContextExtraData);
        }
        /* 在request 放入上下文信息 */
        if (!gatewayContext.getReadRequestData()) {
            exchange.getAttributes().put(GatewayContext.CACHE_GATEWAY_CONTEXT, gatewayContext);
            return chain.filter(exchange);
        }
        /* save gateway context into exchange */
        gatewayContext.getAllRequestData().addAll(request.getQueryParams());
        exchange.getAttributes().put(GatewayContext.CACHE_GATEWAY_CONTEXT, gatewayContext);
        MediaType contentType = headers.getContentType();
        if (headers.getContentLength() > 0) {
            if (MediaType.APPLICATION_JSON.equals(contentType)) {
                return readBody(exchange, chain, gatewayContext);
            }
            if (MediaType.APPLICATION_FORM_URLENCODED.equals(contentType)) {
                return readFormData(exchange, chain, gatewayContext);
            }
        }
        return chain.filter(exchange);

    }


    @Override
    public int getOrder() {
        return FilterOrderEnum.GATEWAY_CONTEXT_FILTER.getOrder();
    }

    /**
     * 因为加入了SQL注入和XSS注入拦截，这里单独判断
     *
     * @return boolean
     */
    private boolean shouldReadRequestInjectionData(ServerWebExchange exchange) {
        return gatewayPluginProperties.isRecordAllLog() || gatewayPluginProperties.shouldSqlInjection(exchange) || gatewayPluginProperties.shouldXssInjection(exchange);
    }


    private Charset getCharsetName(ServerWebExchange exchange) {
        HttpHeaders headers = exchange.getRequest().getHeaders();
        Charset charset = StandardCharsets.UTF_8;
        if (ObjUtil.isNotNull(headers.getContentType())) {
            charset = headers.getContentType().getCharset();
            if (ObjUtil.isNull(charset)) {
                charset = StandardCharsets.UTF_8;
            }
        }
        return charset;
    }

    /**
     * ReadFormData
     * {exchange}
     * {chain}
     *
     * @return {}
     */
    private Mono<Void> readFormData(ServerWebExchange exchange, GatewayFilterChain chain, GatewayContext gatewayContext) {

        return exchange.getFormData().doOnNext(multiValueMap -> {
            gatewayContext.setFormData(multiValueMap);
            gatewayContext.getAllRequestData().addAll(multiValueMap);
        }).then(Mono.defer(() -> {
            Charset charset = getCharsetName(exchange);
            MultiValueMap<String, String> formData = gatewayContext.getFormData();
            if (ObjUtil.isNull(formData) || formData.isEmpty()) {
                return chain.filter(exchange);
            }
            StringBuilder formDataBodyBuilder = new StringBuilder();
            try {
                for (var entry : formData.entrySet()) {
                    var entryKey = entry.getKey();
                    var entryValue = entry.getValue();
                    for (String value : entryValue) {
                        formDataBodyBuilder.append(entryKey).append("=").append(URLEncoder.encode(value, charset.name())).append("&");
                    }
                }
            } catch (UnsupportedEncodingException ignored) {
            }
            /*
             * substring with the last char '&'
             */
            String formDataBodyString = "";
            if (!formDataBodyBuilder.isEmpty()) {
                formDataBodyString = formDataBodyBuilder.substring(0, formDataBodyBuilder.length() - 1);
            }
            /*
             * get data bytes
             */
            byte[] bodyBytes = formDataBodyString.getBytes(charset);
            int contentLength = bodyBytes.length;
            HttpHeaders httpHeaders = new HttpHeaders();
            httpHeaders.putAll(exchange.getRequest().getHeaders());
            httpHeaders.remove(HttpHeaders.CONTENT_LENGTH);
            httpHeaders.setContentLength(contentLength);
            BodyInserter<String, ReactiveHttpOutputMessage> bodyInserter = BodyInserters.fromValue(formDataBodyString);
            CachedBodyOutputMessage cachedBodyOutputMessage = new CachedBodyOutputMessage(exchange, httpHeaders);
            return bodyInserter.insert(cachedBodyOutputMessage, new BodyInserterContext()).then(Mono.defer(() -> {
                ServerHttpRequestDecorator decorator = new ServerHttpRequestDecorator(exchange.getRequest()) {
                    @Override
                    public HttpHeaders getHeaders() {
                        return httpHeaders;
                    }

                    @Override
                    public Flux<DataBuffer> getBody() {
                        return cachedBodyOutputMessage.getBody();
                    }
                };
                return chain.filter(exchange.mutate().request(decorator).build());
            }));
        }));
    }

    /**
     * ReadJsonBody
     * {exchange}
     * {chain}
     *
     * @return {}
     */
    private Mono<Void> readBody(ServerWebExchange exchange, GatewayFilterChain chain, GatewayContext gatewayContext) {
        return DataBufferUtils.join(exchange.getRequest().getBody()).flatMap(dataBuffer -> {
            var cachedFlux = SchemeUtil.getCacheFlux(dataBuffer, exchange, false);
            ServerHttpRequest mutatedRequest = new ServerHttpRequestDecorator(exchange.getRequest()) {
                @Override
                public Flux<DataBuffer> getBody() {
                    return cachedFlux;
                }
            };
            ServerWebExchange mutatedExchange = exchange.mutate().request(mutatedRequest).build();
            return ServerRequest.create(mutatedExchange, MESSAGE_READERS).bodyToMono(String.class).doOnNext(gatewayContext::setRequestBody).then(chain.filter(mutatedExchange));
        });
    }

}